import numpy as np
import torch

# TODO 2.自定义类,构建神经网络模型
class Model_test(torch.nn.Module):
    # 重写__init__方法和forward方法
    def __init__(self,input_num,output_num):
        # 调用父类的构造方法
        super().__init__()
        # 定义网络结构
        self.linear1 = torch.nn.Linear(input_num,3)
        self.linear2 = torch.nn.Linear(3,2)
        self.out = torch.nn.Linear(2,output_num).softmax()

    # 重写forward方法
    def forward(self,x):
        # 加权求和-> 激活函数(默认隐藏层都用relu作为激活函数)
        x = torch.relu(self.linear1(x))
        x = torch.relu(self.linear2(x))
        x = self.out(x)
        return x


if __name__ == '__main__':

    m = Model_test(3, 2)
